import os
import pandas as pd
import numpy as np
import json
import random
import glob
from shutil import copyfile, move
from pathlib import Path

# path = r'/root/project/AutoRepair_T7/data/75/all'

# train_path = r'/root/project/AutoRepair_T7/data/75/train'
# test_path = r'/root/project/AutoRepair_T7/data/75/test'
# val_ratio = 0.15

# for code in os.listdir(path):
#     l = glob.glob(os.path.join(path, code, '*.jpg'))
#     if code in ["TGXID", "TCVPT"]:
#         random.shuffle(l)
#         l = l[:100]
#     train_num = int(len(l) * (1-val_ratio))
#     for i in range(train_num):
#         os.makedirs(os.path.join(train_path, code), exist_ok=True)
#         img_name = Path(l[i]).name
#         copyfile(l[i], os.path.join(train_path, code, img_name))
#         if(code != 'TGXID'):
#             json_name = img_name[:-3] + 'json'
#             copyfile(l[i][:-3]+'json', os.path.join(train_path, code, json_name))
    
#     for i in range(train_num, len(l)):
#         os.makedirs(os.path.join(test_path, code), exist_ok=True)
#         img_name = Path(l[i]).name
#         copyfile(l[i], os.path.join(test_path, code, img_name))
#         if(code != 'TGXID'):
#             json_name = img_name[:-3] + 'json'
#             copyfile(l[i][:-3]+'json', os.path.join(test_path, code, json_name))



# path = r'/data2/autorepair/ruanzhifeng/autorepair_t7_10/t7/T6006/for_label_0315'

# df = pd.DataFrame() 
# df['imgs'] = glob.glob(os.path.join(path, "*/*.jpg"))
# df['code'] = df['imgs'].apply(lambda x: Path(x).parent.name)

# print(df['code'].value_counts())
# print(len(df))

def move_file(img_path, code, memo):
    if code  == 'TGXID':
        return
    print(img_path)
    img_name = Path(img_path).name
    for dp in memo.keys():
        if img_name in memo[dp]:
            os.makedirs(os.path.join(dp, code), exist_ok=True)
            if os.path.exists(os.path.join(dp, 'TGXID', img_name)):
                move(os.path.join(dp, 'TGXID', img_name), os.path.join(dp, code, img_name))
            label_src_path = Path(img_path).with_suffix('.json')
            if os.path.exists(label_src_path):
                label_name = Path(img_path).with_suffix('.json').name
                move(label_src_path, os.path.join(dp, code, label_name))


src_path = r'/data2/autorepair/ruanzhifeng/autorepair_t7_10/t7/t6t7_Charm_TGXID_labeled/T7002'
# dst_paths = [r'/data2/autorepair/ruanzhifeng/tianxinji/data/T6_Vtech/T6009/0621',
#             # r'/data2/autorepair/ruanzhifeng/tianxinji/data/T6_Vtech/T6010/add_data_0705',
#             # r'/data2/autorepair/ruanzhifeng/tianxinji/data/T6_Vtech/T6010/add_data_0712'
            # ]
dst_paths = []
r = r'/data2/autorepair/ruanzhifeng/autorepair_t7_10/t7/T7002'
for folder in os.listdir(r):
    dst_paths.append(os.path.join(r, folder))

memo = {}
for dp in dst_paths:
    imgs = glob.glob(os.path.join(dp, "**/*.jpg"), recursive=True)
    memo[dp] = [Path(x).name for x in imgs]

imgs = glob.glob(os.path.join(src_path, "**/*.jpg"), recursive=True)
df = pd.DataFrame(data={'img_path': imgs})
df['code'] = df['img_path'].apply(lambda x: Path(x).parent.name)

df.apply(lambda x: move_file(img_path=x['img_path'], code=x['code'], memo=memo), axis=1)


# # tmp_imgs = glob.glob(os.path.join(dst_path, "*/*.jpg"))
# df = pd.DataFrame() 
# df['imgs'] = glob.glob(os.path.join(dst_path, "*/*.jpg"))
# df['img_name'] = df['imgs'].apply(lambda x: Path(x).name)
# # df['code'] = df['imgs'].apply(lambda x: Path(x).parent.name)
# ex_imgs = set(df['img_name'].tolist())


# df2 = pd.DataFrame() 
# df2['imgs'] = glob.glob(os.path.join(src_path, "*/*.jpg"))
# df2['img_name'] = df['imgs'].apply(lambda x: Path(x).name)
# df2['code'] = df2['imgs'].apply(lambda x: Path(x).parent.name)

# df2['to_del'] = df2['img_name'].apply(lambda x: x in ex_imgs)

# df2 = df2[df2['to_del']]
# df2.apply(lambda x: os.remove(x['imgs']), axis=1)

# # print(df2['to_del'].value_counts())
